#include <stdio.h>
#include <sys/syscall.h>
#include <immintrin.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <stdbool.h>
#include <string.h>
#include <assert.h>

#include "utils.h"

// trace facilities
uint64_t g_status = 0;

volatile void start_tracing() {
  g_status = 1;
}

volatile void stop_tracing() {
  g_status = 0;
}

/* Initialize tile config */
void init_tile_config (__tilecfg *tileinfo)
{
  int i;
  tileinfo->palette_id = 1;
  tileinfo->start_row = 0;

  for (i = 0; i < 8; ++i)
  {
    tileinfo->colsb[i] = MAX_COLS;
    tileinfo->rows[i] =  MAX_ROWS;
  }

  _tile_loadconfig (tileinfo);
}

/* Set_tiledata_use() - Invoke syscall to set ARCH_SET_STATE_USE */
bool set_tiledata_use()
{
   if (syscall(SYS_arch_prctl, ARCH_REQ_XCOMP_PERM, XFEATURE_XTILEDATA)) 
   {
      printf("\n Fail to do XFEATURE_XTILEDATA \n\n");
      return false;
   }
   else
   {
      printf("\n TILE DATA USE SET - OK \n\n");
      return true;
   }

   return true;
}

void print_matrix_to_file(void* matrix, size_t rows, size_t cols, size_t element_size, const char* filename) {
  FILE* file = fopen(filename, "w");
  if (file == NULL) {
      perror("Failed to open file");
      return;
  }

  if (element_size == sizeof(int8_t)) {
      int8_t* mat = (int8_t*)matrix;
      fprintf(file, "Matrix (int8_t):\n");
      for (size_t i = 0; i < rows; ++i) {
          for (size_t j = 0; j < cols; ++j) {
              fprintf(file, "%4d ", mat[i * cols + j]); // 打印 int8_t 类型的值
          }
          fprintf(file, "\n");
      }
  } else if (element_size == sizeof(int32_t)) {
      int32_t* mat = (int32_t*)matrix;
      fprintf(file, "Matrix (int32_t):\n");
      for (size_t i = 0; i < rows; ++i) {
          for (size_t j = 0; j < cols; ++j) {
              fprintf(file, "%8d ", mat[i * cols + j]); // 打印 int32_t 类型的值
          }
          fprintf(file, "\n");
      }
  } else {
      fprintf(file, "Unsupported element size: %zu\n", element_size);
  }

  fclose(file);
}

void init_matrices(int8_t *A, int8_t *B, int32_t *C, size_t M, size_t N, size_t K) 
{
  for (int i = 0; i < M*K; i++) {
      A[i] = (int8_t)(i % 10 + 1);
      // A[i] = (int8_t)(1);
  }

  for (int i = 0; i < K*N; i++) {
      B[i] = (int8_t)(i % 10 + 1);
      // B[i] = (int8_t)(1);
  }

  for (int i = 0; i < M*N; i++){
        C[i] = 0;
  }
}

bool check_equality(const int32_t* C, const int32_t* expected_C, const size_t M, const size_t N) {
  bool is_correct = true;
  uint32_t err_num = 0;

  for (size_t i = 0; i < M; ++i) {
      for (size_t j = 0; j < N; ++j) {
          if (expected_C[i * N + j] != C[i * N + j]) {
              printf("Mismatch at C[%zu][%zu]: expected %d, got %d\n",
                     i, j, expected_C[i * N + j], C[i * N + j]);
              err_num++;
              if (err_num > 10) {
                  is_correct = false;
                  printf("too many errors, exiting...\n\n");
                  break;
              }
          }
      }
      if (!is_correct) break;
  }

  return is_correct;
}

static void naive_matrix_multiplication(const int8_t *A, const int8_t *B,
                                      const size_t M, const size_t N, const size_t K,
                                      int32_t *expected_C) {
  for (size_t i = 0; i < M; ++i) {
      for (size_t j = 0; j < N; ++j) {
          for (size_t k = 0; k < K; ++k) {
              expected_C[i * N + j] += A[i * K + k] * B[k * N + j];
          }
      }
  }
}

static void matrix_multiplication_with_L1_pack(const int8_t *A, const int8_t *B,
                                             const size_t M, const size_t N, const size_t K,
                                             int32_t *expected_C) {
  for (size_t mb = 0; mb < (M/M_CACHE); mb++) {
      for (size_t nb = 0; nb < (N/N_CACHE); nb++) {
          for (size_t kb = 0; kb < (K/K_CACHE); kb++) {
              const size_t A_tile_offset = (mb*(K/K_CACHE)+kb) * (M_CACHE*K_CACHE) * sizeof(uint8_t);
              const size_t B_tile_offset = (kb*(N/N_CACHE)+nb) * (K_CACHE*N_CACHE) * sizeof(uint8_t); 
              const size_t C_tile_offset = (mb*(N/N_CACHE)+nb) * (M_CACHE*N_CACHE) * sizeof(uint32_t);

              const int8_t* A_tile_ptr = A + A_tile_offset;
              const int8_t* B_tile_ptr = B + B_tile_offset;
              int32_t* C_tile_ptr = (int32_t* )((int8_t *)expected_C + C_tile_offset);

              for (size_t m = 0; m < M_CACHE; m++) {
                  for (size_t n = 0; n < N_CACHE; n++) {
                      for (size_t k = 0; k < K_CACHE; k++) {
                          C_tile_ptr[m*N_CACHE+n] += A_tile_ptr[m*K_CACHE+k] * B_tile_ptr[k*N_CACHE+n];
                      }
                    }
                  }
                }
              }
            }
}

// 矩阵乘法验证函数
void verify_matrix_multiplication(const int8_t* A, const int8_t* B, const int32_t* C, 
                                const size_t M, const size_t N, const size_t K, 
                            const bool pack) {
  // 手动计算矩阵乘法结果
  int32_t* expected_C = (int32_t*)malloc(M * N * sizeof(int32_t));
  memset(expected_C, 0, M * N * sizeof(int32_t));

  // 比较计算结果与程序结果
  if (pack) {
    matrix_multiplication_with_L1_pack(A, B, M, N, K, expected_C);
  } else {
    naive_matrix_multiplication(A, B, M, N, K, expected_C);
  }
  // print_matrix_to_file(expected_C, M, N, 4, "expected_C.txt");

  if (check_equality(C, expected_C, M, N)) {
      printf("Matrix multiplication is correct.\n");
  } else {
      printf("Matrix multiplication is incorrect.\n");
  }

  free(expected_C);
}

void stats(const size_t M, const size_t N, const size_t K) {
    printf("[M_CACHE,N_CACHE,K_CACHE]=[%d,%d,%d]\n", M_CACHE, N_CACHE, K_CACHE);
    printf("[M,N,K]=[%d,%d,%d]\n", M, N, K);
    printf("[TILE_M,TILE_N,TILE_K]=[%d,%d,%d]\n", TILE_M, TILE_N, TILE_K);

    double A_bytes = M*K*sizeof(int8_t);
    double B_bytes = K*N*sizeof(int8_t);
    double C_bytes = M*N*sizeof(int32_t);
    double total_kbytes = (A_bytes+B_bytes+C_bytes) / 1024;

    printf("A KB:     [%.2f] KB\n", A_bytes/1024);
    printf("B KB:     [%.2f] KB\n", B_bytes/1024);
    printf("C KB:     [%.2f] KB\n", C_bytes/1024);
    printf("A+B+C KB:     [%.2f] KB\n", total_kbytes);
}

double get_vnni_avx512_peak_gflops() {
    const double peak_gflops = 2.0 * 64.0 * 2.0 * FREQ; // 64 ops per cycle, 2 FLOPs per op, 2.5 GHz clock

    return peak_gflops;
}

void vnni_avx512_b_layout_transform(int8_t* B, int8_t* Bp, size_t K, size_t N)
{
    assert(K % 4 == 0 && N % 16 == 0);

    for (size_t k = 0; k < K; k += 4)
    {
        for (size_t n = 0; n < N; n += 16)
        {
            for (size_t nn = 0; nn < 16; ++nn)
            {
                for (size_t kk = 0; kk < 4; ++kk)
                {
                    size_t src_idx = (k + kk) * N + (n + nn);           // B[k+kk][n+nn]
                    size_t dst_idx = k * N + n * 4 + nn * 4 + kk;       // pack to [k][n][kk]
                    Bp[dst_idx] = B[src_idx];
                }
            }
        }
    }
}

double get_amx_peak_gflops() {
    const double peak_gflops = 2048 * FREQ;

    return peak_gflops;
}

void amx_b_layout_transform(uint8_t *src, uint8_t *dst, size_t rows_src, size_t cols_src)
{
    for (int k = 0; k < rows_src/4; k++)
    {
        for (int i = 0; i < cols_src; i ++)
        {
            for (int j = 0; j < 4; j++)
            {
                dst[k*cols_src*4 + i*4 + j] = src[k*cols_src*4 + j*cols_src + i];
            }
        }
    }
}

void amx_b_layout_transform_with_pack(uint8_t *src, uint8_t *dst, size_t rows_src, size_t cols_src)
{
    for (size_t kb = 0; kb < (rows_src/K_CACHE); kb++) {
        for (size_t nb = 0; nb < (cols_src/N_CACHE); nb++) {
            const size_t B_tile_offset = (kb*(cols_src/N_CACHE)+nb) * (K_CACHE*N_CACHE) * sizeof(uint8_t); 
  
                const int8_t* src_tile_ptr = src + B_tile_offset;
                int8_t* dst_tile_ptr = dst + B_tile_offset;
  
                for (size_t k = 0; k < K_CACHE/4; k++) {
                    for (size_t n = 0; n < N_CACHE; n++) {
                        for (int p = 0; p < 4; ++p) {
                            dst_tile_ptr[k*N_CACHE*4+n*4+p] = src_tile_ptr[k*N_CACHE*4 + p*N_CACHE + n];
                        }
                    }
                }
            }
        }
}